import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from collections import Counter

# 1. 加载 Iris 数据集
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/iris/iris.data"
columns = ['SepalLength', 'SepalWidth', 'PetalLength', 'PetalWidth', 'Class']
df = pd.read_csv(url, header=None, names=columns)

# 显示数据的前几行
print(df.head())

# 2. 数据预处理
# 将数据分为特征（X）和标签（y）
X = df.drop('Class', axis=1).values
y = df['Class'].values


# 3. 实现KNN算法
def euclidean_distance(x1, x2):
    """计算两个点之间的欧几里得距离"""
    return np.sqrt(np.sum((x1 - x2) ** 2))


def knn(X_train, y_train, X_test, k=3):
    """实现KNN分类算法"""
    predictions = []

    # 对每个测试点进行分类
    for test_point in X_test:
        # 计算训练集所有点与测试点之间的距离
        distances = [euclidean_distance(test_point, train_point) for train_point in X_train]

        # 找到距离最小的k个点
        k_indices = np.argsort(distances)[:k]
        k_nearest_labels = [y_train[i] for i in k_indices]

        # 统计k个邻居中最多的标签
        most_common = Counter(k_nearest_labels).most_common(1)
        predictions.append(most_common[0][0])

    return np.array(predictions)


# 4. 划分训练集和测试集
# 这里我们使用80%的数据作为训练集，20%的数据作为测试集
train_size = int(0.8 * len(df))
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]

# 5. 使用KNN进行预测
k = 3
y_pred = knn(X_train, y_train, X_test, k)

# 6. 评估模型的准确性
accuracy = np.mean(y_pred == y_test)
print(f"Accuracy of KNN (k={k}): {accuracy * 100:.2f}%")

# 7. 可视化数据和分类结果
# 将前两个特征进行可视化（即：SepalLength 和 SepalWidth）
plt.figure(figsize=(8, 6))
plt.scatter(X_test[:, 0], X_test[:, 1],
            c=[{'Iris-setosa': 0, 'Iris-versicolor': 1, 'Iris-virginica': 2}[label] for label in y_pred],
            cmap='viridis', label="Test Points")
plt.scatter(X_train[:, 0], X_train[:, 1],
            c=[{'Iris-setosa': 0, 'Iris-versicolor': 1, 'Iris-virginica': 2}[label] for label in y_train],
            cmap='coolwarm', label="Train Points", alpha=0.5)
plt.title('KNN Classification of Iris Dataset')
plt.xlabel('Sepal Length')
plt.ylabel('Sepal Width')
plt.legend()
plt.show()
